import numpy as np
import socket
import json
from scipy.signal import butter, filtfilt, spectrogram
from sklearn.cross_decomposition import CCA
import time
# 导入新创建的RDA接口模块
from cca.rda_interface_pic import RDAInterface
# 新 时频能量图
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 全局绘图变量
fig = None
ax = None
mesh = None
channels_to_plot = ['P3', 'Pz', 'P4', 'P8', 'O1', 'O2']


def load_control_config():
    """加载控制配置文件"""
    with open('ControlCode_config.json', 'r') as f:
        return json.load(f)


def load_sliding_window_config():
    """加载滑动窗口配置参数"""
    return {
        'window_length': 1.0,  # 窗长（秒）
        'step_size': 0.25,  # 步长（秒）
        'correlation_threshold': 0.3  # 相关系数阈值
    }


def create_tcp_client():
    """创建TCP客户端连接"""
    client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    return client


def connect_to_drone(client, port):
    """连接到无人机控制程序"""
    try:
        client.connect(('192.168.43.52', port))
        print("成功连接到Drone Flying程序")
        return True
    except ConnectionRefusedError:
        print("无法连接到无人机控制程序。请确保程序正在运行并监听端口。")
        return False


def send_command(client, command):
    """发送控制命令"""
    ascii_code = int(command, 16)  # 将十六进制字符串转换为整数
    command_char = chr(ascii_code)  # 将ASCII码转换为对应字符
    print(f"发送ASCII码: 0x{command} ({ascii_code}) -> 字符: {command_char}")
    try:
        client.send(command_char.encode())
        time.sleep(0.1)  # 添加短暂延时确保命令被正确处理
    except BrokenPipeError:
        print("TCP连接已断开，无法发送命令。")
        return False
    return True


def butter_bandpass_filter(data, lowcut, highcut, sample_rate, order):
    """带通滤波器"""
    nyq = 0.5 * sample_rate
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    y = filtfilt(b, a, data)
    return y


def get_cca_reference_signals(data_len, target_freq, sampling_rate):
    """生成CCA参考信号"""
    reference_signals = []
    t = np.arange(0, (data_len / (sampling_rate)), step=1.0 / (sampling_rate))

    reference_signals.append(np.sin(np.pi * 2 * target_freq * t))
    reference_signals.append(np.cos(np.pi * 2 * target_freq * t))
    reference_signals.append(np.sin(np.pi * 4 * target_freq * t))
    reference_signals.append(np.cos(np.pi * 4 * target_freq * t))

    return np.array(reference_signals)


def find_correlation(n_components, eeg_data, reference_signals):
    """使用CCA计算相关性"""
    cca = CCA(n_components)
    cca.fit(eeg_data.T, reference_signals.T)
    O1_a, O1_b = cca.transform(eeg_data.T, reference_signals.T)

    corr = np.zeros(n_components)
    for i in range(n_components):
        corr[i] = np.corrcoef(O1_a[:, i], O1_b[:, i])[0, 1]

    return np.max(corr)


def setup_time_frequency_plot(sfreq):
    """初始化动态时频图"""
    global fig, ax, mesh

    # 开启交互模式
    plt.ion()

    fig, ax = plt.subplots(figsize=(12, 8))

    # 初始化一个空图
    nperseg = int(2 * sfreq)
    t = np.arange(nperseg) / sfreq
    f = np.linspace(5, 20, 100)  # 预设频率轴

    # 使用一个零矩阵初始化pcolormesh
    mesh = ax.pcolormesh(t, f, np.zeros((100, len(t))), shading='gouraud')

    ax.set_title(f'时频能量图 - 通道: {channels_to_plot}')
    ax.set_ylabel('频率 [Hz]')
    ax.set_xlabel('时间 [秒]')
    ax.set_ylim([5, 20])
    fig.colorbar(mesh, ax=ax, label='能量/频率 [dB/Hz]')

    plt.show(block=False)  # 非阻塞显示
    plt.pause(0.1)  # 短暂暂停以确保图表已创建

    return fig, ax, mesh

def update_time_frequency_plot(eeg_data, sfreq, ch_names, fig, ax, mesh):
    """更新动态时频图的数据"""
    # 查找要绘制的通道的索引
    channel_indices = [ch_names.index(ch) for ch in channels_to_plot if ch in ch_names]

    if not channel_indices:
        print("警告: 指定的绘图通道未找到。")
        return

    selected_data = eeg_data[channel_indices, :]
    averaged_data = np.mean(selected_data, axis=0)

    nperseg = int(2 * sfreq)
    noverlap = int(1 * sfreq)

    if len(averaged_data) < nperseg:
        print(f"警告: 平均数据长度 ({len(averaged_data)}) 小于STFT窗口长度 ({nperseg})。跳过绘图。")
        return

    # 对输入的信号进行短时傅里叶变换（STFT）
    '''
    f：一个一维 NumPy 数组，包含了时频图的频率轴。
    t：一个一维 NumPy 数组，包含了时频图的时间轴。
    Sxx：一个二维 NumPy 数组，是 spectrogram 函数的核心输出。这个矩阵包含了信号在每个时间和频率点上的能量或功率谱密度。这个矩阵的数据最终会被 plt.pcolormesh 函数用来绘制热力图。
    '''
    f, t, Sxx = spectrogram(averaged_data, sfreq, nperseg=nperseg, noverlap=noverlap)

    # 更新绘图数据
    mesh.set_array(10 * np.log10(Sxx).ravel())
    mesh.set_clim(np.min(10 * np.log10(Sxx)), np.max(10 * np.log10(Sxx)))

    # 调整x轴范围以匹配最新数据
    ax.set_xlim(t[0], t[-1])

    # 刷新图表
    fig.canvas.draw()
    fig.canvas.flush_events()

def process_eeg_data(raw_data, sfreq, selected_channels, ch_names):
    """处理EEG数据并返回预测结果和最大相关系数

    参数:
    raw_data - EEG数据矩阵 (n_channels, n_times)
    sfreq - 采样率
    selected_channels - 选定用于分析的通道列表
    ch_names - 所有通道的名称列表
    """
    # 滤波参数
    lowcut = 6
    highcut = 80
    filter_order = 4

    # SSVEP频率设置
    ssvep_freqs = np.array([12.0, 15.0, 13.0, 7.0, 10, 11.0])  # 左、前、右、后、上、下

    # 对选定的通道数据进行带通滤波
    filtered_data = np.zeros((len(selected_channels), raw_data.shape[1]))
    for i, ch_name in enumerate(selected_channels):
        # 找到通道索引
        try:
            ch_idx = ch_names.index(ch_name)
            filtered_data[i, :] = butter_bandpass_filter(
                raw_data[ch_idx, :],
                lowcut, highcut, sfreq, filter_order
            )
        except ValueError:
            print(f"警告: 未找到通道 {ch_name}，使用默认值")
            filtered_data[i, :] = np.zeros(raw_data.shape[1])

    # 计算每个频率的相关性
    correlations = np.zeros(len(ssvep_freqs))
    for freq_idx, freq in enumerate(ssvep_freqs):
        reference = get_cca_reference_signals(raw_data.shape[1], freq, sfreq)
        correlations[freq_idx] = find_correlation(1, filtered_data, reference)

    # 返回相关性最大的频率索引和对应的相关系数
    max_corr_idx = np.argmax(correlations)
    return max_corr_idx, correlations[max_corr_idx]


def realtime_data_callback(data_buffer, sfreq, ch_names, markers, client, config, selected_channels,
                           correlation_threshold, freq_to_command, fig, ax, mesh):
    """实时数据回调处理函数"""
    try:
        # 输出当前窗口信息
        window_duration = data_buffer.shape[1] / sfreq
        print(f"处理窗口: {data_buffer.shape[1]}个样本点, 时长: {window_duration:.2f}秒, 采样率: {sfreq}Hz")

        # 处理当前缓冲区的数据
        prediction, max_correlation = process_eeg_data(data_buffer, sfreq, selected_channels, ch_names)
        print(f"预测结果: {prediction}, 相关系数: {max_correlation:.3f}")

        # 动态更新时频图
        update_time_frequency_plot(data_buffer, sfreq, ch_names, fig, ax, mesh)

        # 只有当相关系数超过阈值时才执行命令
        if max_correlation >= correlation_threshold:
            command_type = freq_to_command.get(prediction)
            if command_type:
                command = config[command_type]
                print(f"检测到的动作: {command_type}, 相关系数: {max_correlation:.3f}")
                send_command(client, command)
        else:
            print(f"相关系数 ({max_correlation:.3f}) 低于阈值，丢弃该窗口结果")

        # 处理标记
        if markers:
            for marker in markers:
                print(f"标记: {marker.description}")

        print("-" * 50)  # 分隔线，便于查看每次处理结果

    except Exception as e:
        print(f"实时数据处理错误: {e}")


def process_realtime_mode(client, config, selected_channels):
    """处理实时模式，使用RDAInterface类接收RDA数据流"""
    # 加载滑动窗口配置
    window_config = load_sliding_window_config()
    window_length = window_config['window_length']
    step_size = window_config['step_size']
    correlation_threshold = window_config['correlation_threshold']

    print(f"滑动窗口配置: 窗长={window_length}秒, 步长={step_size}秒, 相关阈值={correlation_threshold}")

    # 频率到控制命令的映射
    freq_to_command = {
        0: "Left_Code",  # 7.5Hz - 左 ✓ 正确
        1: "Forward_Code",  # 5Hz  - 前 ✓ 正确
        2: "Right_Code",  # 10Hz - 右 ✓ 正确
        3: "Back_Code",  # 6.5Hz - 后
        4: "Up_Code",  # 12Hz - 上 ✓ 正确
        5: "Down_Code"  # 15Hz - 下 ✓ 正确
    }

    # 设置RDA连接参数
    host = "localhost"
    port = int("51244")
    buffer_size = 4  # 使用窗长作为缓冲区大小

    # 创建RDA接口实例
    rda_interface = RDAInterface(host=host, port=port, buffer_size=buffer_size)

    # 滑动窗口处理相关变量
    last_process_time = 0

    # 初始化绘图
    fig, ax, mesh = setup_time_frequency_plot(sfreq=1000)  # 假设采样率为1000

    try:
        # 连接到RDA服务器
        if rda_interface.connect():
            print(f"已连接到RDA服务器: {host}:{port}")
            print(f"窗长: {window_length}秒, 步长: {step_size}秒")

            # 创建回调函数的闭包，传入必要的参数
            def callback_wrapper(data, sfreq, ch_names, markers):
                nonlocal last_process_time
                current_time = time.time()

                # 检查是否达到步长间隔
                if current_time - last_process_time >= step_size:
                    realtime_data_callback(data, sfreq, ch_names, markers,
                                           client, config, selected_channels,
                                           correlation_threshold, freq_to_command,
                                           fig, ax, mesh)
                    last_process_time = current_time

            # 设置回调函数
            rda_interface.set_callback(callback_wrapper)

            # 开始处理数据
            if rda_interface.start_processing():
                print("正在接收和处理实时数据流，按Ctrl+C停止...")
                try:
                    # 保持程序运行直到用户中断
                    while rda_interface.is_connected():
                        plt.pause(0.1)  # 在循环中添加暂停以允许GUI更新
                except KeyboardInterrupt:
                    print("用户中断接收")
                finally:
                    rda_interface.stop_processing()
                    print("已停止实时数据处理")
            else:
                print("启动RDA数据处理失败")
        else:
            print("无法连接到RDA服务器")
    except Exception as e:
        print(f"实时模式错误: {e}")


def main():
    # 加载控制配置
    config = load_control_config()
    port = config['listenPort']

    # 创建TCP客户端并连接
    client = create_tcp_client()
    if not connect_to_drone(client, port):
        return

    # 设置EEG数据采集参数
    selected_channels = ['P3', 'Pz', 'P4', 'P8', 'O1', 'O2']

    # 启动实时数据处理模式
    try:
        print("启动实时数据处理模式...")
        process_realtime_mode(client, config, selected_channels)
    except Exception as e:
        print(f"处理错误: {e}")
    finally:
        # 关闭连接
        client.close()
        print("TCP连接已关闭")
        plt.close('all')  # 关闭所有matplotlib窗口


if __name__ == "__main__":
    main()